import numpy as np
import matplotlib.pyplot as plt

plt.style.use("ggplot")


# 取 omega_R 作为单位，定义为 1， 不妨设 omegaR = 200 Hz
# 则时间的基本单位为 1/omegaR
# 若取 T = 1s, 在 omega_R 这个单位下，T = omegaR 是一个较大值
omegaR = 200
T = omegaR
Delta = np.linspace(-3.5,3.5,10000)
p = (1-np.cos(T*Delta))/2/(1+Delta**2)
fig,ax = plt.subplots()
ax.plot(Delta,p)
ax.set_title(r"Zeroth-Order Approximation: $P^{(0)}_{1\rightarrow2}(2t_0+T)$",fontsize=23)
ax.set_xlabel(r"$\Delta/\Omega_R,\Omega_R = 200 Hz,T = 1s$",fontsize=20)
ax.set_ylabel(r"Transition Probability",fontsize=20)
fig.set_size_inches((12.6,9))
fig.savefig("RamseyFringes_0order.png",dpi=400)


omegaR = 200
T = omegaR
Delta = np.linspace(-3.8,3.8,10000)
omega = np.sqrt(1+Delta**2)
t0 = np.pi/2
p = (2*np.sin(omega*t0/2)/omega*(np.cos(omega*t0/2)*np.cos(Delta*T/2) - Delta/omega*np.sin(omega*t0/2)*np.sin(Delta*T/2)) )**2
#p = np.sin(omega*t0/2)**2/omega**2*( (1-np.cos(omega*t0))*(1-np.cos(Delta*T)) + 2*Delta/omega*np.sin(omega*t0)*np.sin(Delta*T) + Delta**2/omega**2*(1+np.cos(omega*t0))*(1+np.cos(Delta*T)))

fig,ax = plt.subplots()
ax.plot(Delta,p)
ax.set_title(r"Precise Result: $P_{1\rightarrow2}(2t_0+T)$",fontsize=23)
ax.set_xlabel(r"$\Delta/\Omega_R,\Omega_R = 200 Hz,T = 1s$",fontsize=20)
ax.set_ylabel(r"Transition Probability",fontsize=20)
fig.set_size_inches((12.6,9))
fig.savefig("RamseyFringes_precise.png",dpi=400)



